#include "consensus.hpp"
#include "snp_call.hpp"
#include "cartesian_product.hpp"
#include "option_parse.hpp"


void change_col_genome(prog_options &options, color_call_vector_t &snp_calls, std::map<std::string, std::string> &genome_seqs)
{
    std::vector< std::vector<double> > M;

    for(unsigned int i = 1; i < snp_calls.size(); i++)
    {
        ColorSnpCall &snp = snp_calls[i];

        int max_index = solve_col_snp(options, snp, M);
    
        // no sequences spanned the snp block.
        if( max_index < 0)
            continue;

        for(unsigned int j = 0; j < snp.end - snp.start; j++)
            genome_seqs[snp.chr][snp.start + j] = snp.blocks[max_index][j].base;
    }
}

void change_seq_genome(prog_options &options, seq_call_vector_t &snp_calls, std::map<std::string, std::string> &genome_seqs)
{
    error("change_seq_genome function is not implemented!");
    for(unsigned int i = 0; i < snp_calls.size(); i++)
    {
        SeqSnpCall &snp = snp_calls[i];
    
        // should we use some other method here?
        if( snp.seq.empty() )
            continue;
        
        int max_index = std::max_element(snp.probs.begin(), snp.probs.end()) - snp.probs.begin();

        genome_seqs[snp.chr][snp.start] = snp.blocks[max_index].base;
    }
}

void do_consensus(int argc, char **argv)
{
    prog_options options = parse_consensus_options(argc, argv);
    std::map<std::string, std::string> genome_seqs;

    parse_fasta(options.fasta_file.c_str(), genome_seqs);

    std::ifstream is(options.snp_file.c_str(), std::ios_base::in);

    if(!is)
        error("Could not read file " + options.snp_file);

    unsigned int ftype, fversion;
    boost::iostreams::filtering_istream in;
    in.push(boost::iostreams::gzip_decompressor());
    in.push(is);
    boost::archive::binary_iarchive ia(in);
    ia >> ftype;
    ia >> fversion;
    if(ftype == SNP_CALL_COL_FILE_TYPE)
    {
        if(fversion != SNP_CALL_COL_FILE_VERSION)
            error("bad version for color snpcall file.");
        color_call_vector_t snp_cands;
        ia >> options.diploid;
        ia >> snp_cands;

        multiply_col_snp_calls_by_priors(options, snp_cands);
        change_col_genome(options, snp_cands, genome_seqs);
    }
    else if(ftype == SNP_CALL_SEQ_FILE_TYPE)
    {
        if(fversion != SNP_CALL_SEQ_FILE_VERSION)
            error("bad version for sequence snpcall file.");
        seq_call_vector_t cv;
        ia >> options.diploid;
        ia >> cv;

        multiply_seq_snp_calls_by_priors(options, cv);
        
        for(unsigned int i = 1; i < cv.size(); i++)
        {
            SeqSnpCall &snp = cv[i];

            if( snp.seq.empty() )
                continue;

            unsigned int max_idx = std::max_element(snp.probs.begin(), snp.probs.end()) - snp.probs.begin();
            genome_seqs[snp.chr][snp.start] = snp.blocks[max_idx].base;
        }
    }

    write_fasta(std::cout, genome_seqs);
}
